import csv

# 小球参数
d = 0.05  # 直径 m
c = 1638.7  # 比热容 J/(kg·K)
k = 0.153  # 导热系数 W/(m·K)
m = 45 / 1000  # 质量 kg
rho = m / ((4 / 3) * 3.1414926 * (d / 2) ** 3)  # 密度 kg/m3
t0 = 298  # 初始温度 K
s = 0  # 内热源 W/m3

# 环境参数
tf = 1073  # 环境温度 K
h = 8.0897  # 对流传热系数 W/(m2·K)


def explicit(start: list, x: float, y: float) -> list:
    """
    时间显式
    :param start: 初值或迭代后的值
    :param x: 点长度
    :param y: 时间步长
    :return: 迭代后的值
    """
    end = []

    # 计算第一个节点 对称绝热条件
    ap1, aw0, ae0, ap0 = rho * c * x / y, 0, k / x, (rho * c * x / y) - (k / x)  # 系数
    tw0, te0, tp0 = 0, start[1], start[0]
    b = s * x
    tp1 = (aw0 * tw0 + ae0 * te0 + ap0 * tp0 + b) / ap1
    end.append(tp1)

    # 计算中间节点
    ap1 = rho * c * x / y
    aw0 = k / x
    ae0 = k / x
    ap0 = (rho * c * x / y) - (k / x) - (k / x)
    b = s * x
    for i in range(1, len(start) - 1):
        tw0 = start[i - 1]
        te0 = start[i + 1]
        tp0 = start[i]
        tp1 = (aw0 * tw0 + ae0 * te0 + ap0 * tp0 + b) / ap1
        end.append(tp1)

    # 计算最后一个节点 第三类边界条件
    ap1 = rho * c * x / y
    aw0 = k / x
    ae0 = 0
    ap0 = (rho * c * x / y) - (k / x) - (1 / ((1 / h) + (x / k)))
    tw0 = start[-2]
    te0 = 0
    tp0 = start[-1]
    b = s * x + (tf / ((1 / h) + (x / (2 * k))))
    tp1 = (aw0 * tw0 + ae0 * te0 + ap0 * tp0 + b) / ap1
    end.append(tp1)

    return end


def implicit_tdma(start: list, x: float, y: float):
    """
    时间隐式 TDMA求解
    :param start: 初值或迭代后的值
    :param x: 节点长度
    :param y: 时间步长
    :return: 迭代后的值
    """
    end, p, q = [0] * len(start), [], []

    # 计算第一个节点pq
    temp_p = (k / x) / ((rho * c * x / y) + (k / x))
    temp_q = (((rho * c * x / y) * start[0]) + s * x) / ((rho * c * x / y) + (k / x))
    p.append(temp_p)
    q.append(temp_q)

    # 计算中间节点pq
    for i in range(1, len(start) - 1):
        b = s * x
        aw1 = k / x
        ap1 = (rho * c * x / y) + k / x + k / x
        ae1 = k / x
        ap0 = rho * c * x / y
        tp0 = start[i]
        temp_p = ae1 / (ap1 - aw1 * p[-1])
        temp_q = (ap0 * tp0 + b + aw1 * q[-1]) / (ap1 - aw1 * p[-1])
        p.append(temp_p)
        q.append(temp_q)

    # 计算最后一个节点温度
    b = s * x + (tf / ((1 / h) + (x / (2 * k))))
    ap1 = (rho * c * x / y) + (k / x) + (1 / (1 / ((1 / h) + (x / (2 * k)))))
    ap0 = rho * c * x / y
    tp0 = start[-1]
    aw1 = k / x
    end[-1] = (ap0 * tp0 + b + aw1 * q[-1]) / (ap1 - aw1 * p[-1])

    # 计算中间节点和第一个节点温度
    for i in range(len(start) - 2, -1, -1):
        end[i] = p[i] * end[i + 1] + q[i]

    return end


def implicit_fp(start: list, x: float, y: float, z: float, max_step: int) -> list:
    """
    时间隐式 固定点迭代 Fixed Point Iteration
    :param start: 初值或迭代后的值
    :param x: 节点长度
    :param y: 时间步长
    :param z: 残差要求
    :param max_step: 最大迭代步数
    :return: 每一步迭代的值 二维数组
    """
    end = [[0] + start]
    iteration_start = start
    # start用来查找起始时间步的值
    # inside_iteration_start用来更新迭代值
    # 注意区分下面用什么值

    for i in range(max_step):
        temp = start
        flag = True

        # 计算第一个节点迭代值
        b = s * x
        ap1 = (rho * c * x / y) + (k / x)
        ae1 = k / x
        ap0 = (rho * c * x / y)
        temp[0] = (ae1 * iteration_start[1] + ap0 * start[0] + b) / ap1
        if temp[0] - iteration_start[0] >= z:
            flag = False

        # 计算中间节点迭代值
        b = s * x
        aw1 = k / x
        ap1 = (rho * c * x / y) + (k / x) + (k / x)
        ae1 = k / x
        ap0 = rho * c * x / y
        for ii in range(1, len(start) - 2):
            temp[i] = (aw1 * iteration_start[i - 1] + ae1 * iteration_start[i + 1] + ap0 * start[i] + b) / ap1
            if temp[i] - iteration_start[i] >= z:
                flag = False

        # 计算最后一个节点迭代值
        b = tf / ((1 / h) + (x / k))
        aw1 = k / x
        ap0 = rho * c * x / y
        ap1 = (rho * c * x / y) + (k / x) + (1 / ((1 / h) + (x / k)))
        temp[-1] = (aw1 * iteration_start[-2] + ap0 * start[-1] + b) / ap1
        if temp[-1] - iteration_start[-1] >= z:
            flag = False

        iteration_start = temp  # 用计算完的代替初值重新迭代
        end.append([i + 1] + iteration_start)

        if flag:  # 残差符合要求时直接收敛
            break

    return end


def implicit_gs(start: list, x: float, y: float) -> list:
    """
    时间隐式 G-S迭代
    :param start: 初值或迭代后的值
    :param x: 节点长度
    :param y: 时间步长
    :return: 每一步迭代的值 二维数组
    """
    end = []
    return end


def implicit_adi(start: list, x: float, y: float) -> list:
    """
    时间隐式 ADI迭代
    :param start: 初值或迭代后的值
    :param x: 节点长度
    :param y: 时间步长
    :return: 每一步迭代的值 二维数组
    """
    end = []
    return end


def cal_residual(pre: list, cur: list) -> list:
    """
    计算残差
    :param pre: 上一次迭代的值
    :param cur: 本次迭代的值
    :return: 残差
    """
    res = []
    for i in range(1, len(pre)):
        res.append(cur[i])
        res.append(abs(cur[i] - pre[i]))
    return res


def make_header(num: int, met: int) -> list:
    """
    制作表头
    :param num: 节点个数
    :param met: 采用方法
    :return: 表头
    """
    header = ['Time']
    if met == 1 or met == 21:
        for j in range(num):
            temp = 'Temp' + str(j + 1)
            header.append(temp)
    else:
        header.append('Iteration Step')
        for j in range(num):
            temp = 'Temp' + str(j + 1)
            header.append(temp)
            temp = 'Residual' + str(j + 1)
            header.append(temp)
    return header


if __name__ == '__main__':
    node = 50  # 离散节点个数
    node_length = d / (2 * node)  # 每个节点的长度
    time_step = 0.1  # 时间步长
    total_step = 200  # 总共计算步数
    residual = 0.01  # 残差 在显式和隐式TDMA中没用
    max_iteration_step = 50  # 最大迭代步数

    '''
    1 时间显式
    21 时间隐式 TDMA
    22 时间隐式 FP
    23 时间隐式 GS
    24 时间隐式 ADI
    '''
    method = 22

    if method == 1:
        file = open('result_explicit.csv', 'w', newline='')
        writer = csv.writer(file)
        writer.writerow(make_header(node, method))  # 写表头

        data = [t0] * node
        for a in range(total_step):
            time = [time_step * (a + 1)]
            new_data = explicit(data, node_length, time_step)
            data = new_data
            writer.writerow(time + new_data)

        file.close()

    elif method == 21:
        file = open('result_implicit_tdma.csv', 'w', newline='')
        writer = csv.writer(file)
        writer.writerow(make_header(node, method))  # 写表头

        data = [t0] * node
        for a in range(total_step):
            time = [time_step * (a + 1)]
            new_data = implicit_tdma(data, node_length, time_step)
            data = new_data
            writer.writerow(time + new_data)

        file.close()

    elif method == 22:
        file = open('result_implicit_fp.csv', 'w', newline='')
        writer = csv.writer(file)
        writer.writerow(make_header(node, method))  # 写表头

        data = [t0] * node
        for a in range(total_step):
            time = [time_step * (a + 1)]
            new_data = implicit_fp(data, node_length, time_step, residual, max_iteration_step)
            data = new_data[-1]
            if a == 0 or (a + 1) % 100 == 0:  # 每100步保存一次
                for aa in range(len(new_data)):
                    writer.writerow(new_data[aa])
                    # if aa == 0:
                    #     writer.writerow(time + [new_data[aa][0]] + cal_residual(new_data[aa], new_data[aa]))
                    # else:
                    #     writer.writerow([' '] + [new_data[aa][0]] + cal_residual(new_data[aa - 1], new_data[aa]))

        file.close()

    elif method == 23:
        file = open('result_implicit_gs.csv', 'w', newline='')
        writer = csv.writer(file)
        writer.writerow(make_header(node, method))  # 写表头

        data = [t0] * node
        for a in range(total_step):
            time = [time_step * (a + 1)]
            new_data = implicit_gs(data, node_length, time_step)
            data = new_data
            writer.writerow(time + new_data)

        file.close()

    elif method == 24:
        file = open('result_implicit_adi.csv', 'w', newline='')
        writer = csv.writer(file)
        writer.writerow(make_header(node, method))  # 写表头

        data = [t0] * node
        for a in range(total_step):
            time = [time_step * (a + 1)]
            new_data = implicit_adi(data, node_length, time_step)
            data = new_data
            writer.writerow(time + new_data)

        file.close()

    else:
        print('There is NO match method!')
